from requests import Session
from src.database.manager import DBManager
from src.database.models import VideoVectorMapping
from src.utils.logger import log_info


class VideoVectorMappingDAO:
    """Data Access Object for VideoVectorMapping."""
    def __init__(self, session: Session = DBManager().get_session()):
        self.session = session
    
    def add(self, video_id: int, faiss_id: int, vector_type: str, frame_number: int, vector_dimension: int) -> VideoVectorMapping:
        # Check if a record with the same video_id, vector_type, and frame_number exists
        log_info("Adding video vector mapping", ext={
            "video_id": video_id,
            "faiss_id": faiss_id,
            "vector_type": vector_type,
            "frame_number": frame_number,
            "vector_dimension": vector_dimension
        })
        existing_mapping = self.session.query(VideoVectorMapping).filter_by(
            video_id=video_id, 
            vector_type=vector_type, 
            frame_number=frame_number
        ).first()
        
        if existing_mapping:
            # Update the existing record
            existing_mapping.faiss_id = faiss_id
            existing_mapping.vector_dimension = vector_dimension
            self.session.commit()
            return existing_mapping
        else:
            # Add a new record
            mapping = VideoVectorMapping(
                video_id=video_id, 
                faiss_id=faiss_id, 
                vector_type=vector_type,
                frame_number=frame_number, 
                vector_dimension=vector_dimension
            )
            self.session.add(mapping)
            self.session.commit()
            return mapping

    def update(self, faiss_id: int, video_id: int = None, vector_type: str = None, 
                frame_number: int = None, vector_dimension: int = None) -> VideoVectorMapping:
        mapping = self.get_by_faiss_id(faiss_id)
        if not mapping:
            return None
            
        if video_id is not None:
            mapping.video_id = video_id
        if vector_type is not None:
            mapping.vector_type = vector_type
        if frame_number is not None:
            mapping.frame_number = frame_number
        if vector_dimension is not None:
            mapping.vector_dimension = vector_dimension
            
        self.session.commit()
        return mapping

    def delete(self, video_vector_mapping_id: int) -> bool:
        mapping = self.get_by_faiss_id(video_vector_mapping_id)
        if not mapping:
            return False
            
        self.session.delete(mapping)
        self.session.commit()
        return True

    def get_video_id_by_faiss_id(self, faiss_id: int) -> int:
        """
        Get video_id by faiss_id.
        
        Args:
            faiss_id: The faiss_id to look up
            
        Returns:
            The corresponding video_id if found, or None if not found
        """
        mapping = self.session.query(VideoVectorMapping).filter_by(faiss_id=faiss_id).first()
        if mapping:
            return mapping.video_id
        return None

    def get_by_faiss_id(self, faiss_id: int) -> VideoVectorMapping:
        """
        Get VideoVectorMapping by faiss_id.
        
        Args:
            faiss_id: The faiss_id to look up
            
        Returns:
            The corresponding VideoVectorMapping if found, or None if not found
        """
        return self.session.query(VideoVectorMapping).filter_by(faiss_id=faiss_id).first()
    
    def truncate(self):
        """
        Truncate the VideoVectorMapping table.
        """
        self.session.query(VideoVectorMapping).delete()
        self.session.commit()